import importlib.metadata
import platform
import torch
import sys
import os
import tensorflow as tf

def safe_print(msg):
    try:
        print(msg)
    except Exception as e:
        print(f"打印信息失败: {e}")

def check_pytorch():
    try:
        torch_version = importlib.metadata.version('torch')
        safe_print(f"PyTorch 版本: {torch_version}")
        safe_print(f"CUDA 可用: {torch.cuda.is_available()}")
        safe_print(f"CUDA 版本: {torch.version.cuda if hasattr(torch.version, 'cuda') else 'N/A'}")
        safe_print(f"GPU 数量: {torch.cuda.device_count()}")
        if torch.cuda.is_available():
            safe_print(f"当前 GPU: {torch.cuda.get_device_name(0)}")
    except importlib.metadata.PackageNotFoundError:
        safe_print("PyTorch 未安装")
    except Exception as e:
        safe_print(f"检测 PyTorch 时发生异常: {e}")

def check_lightgbm():
    try:
        lgb_version = importlib.metadata.version('lightgbm')
        safe_print(f"\nLightGBM 版本: {lgb_version}")
        try:
            import lightgbm as lgb
            try:
                params = {'device': 'gpu', 'gpu_platform_id': 0, 'gpu_device_id': 0}
                train_data = lgb.Dataset([[1]], label=[0])
                lgb.train(params, train_data, num_boost_round=1)
                safe_print("LightGBM GPU支持: 可用")
            except Exception as e:
                safe_print(f"LightGBM GPU支持: 不可用 ({str(e)})")
        except Exception as e:
            safe_print(f"导入 lightgbm 失败: {e}")
    except importlib.metadata.PackageNotFoundError:
        safe_print("\nLightGBM 未安装")
    except Exception as e:
        safe_print(f"检测 LightGBM 时发生异常: {e}")

def check_catboost():
    try:
        cb_version = importlib.metadata.version('catboost')
        safe_print(f"\nCatBoost 版本: {cb_version}")
        try:
            from catboost import CatBoostClassifier
            try:
                model = CatBoostClassifier(task_type='GPU')
                safe_print("CatBoost GPU支持: 可用")
            except Exception as e:
                safe_print(f"CatBoost GPU支持: 不可用 ({str(e)})")
        except Exception as e:
            safe_print(f"导入 catboost 失败: {e}")
    except importlib.metadata.PackageNotFoundError:
        safe_print("\nCatBoost 未安装")
    except Exception as e:
        safe_print(f"检测 CatBoost 时发生异常: {e}")

def check_xgboost():
    try:
        xgb_version = importlib.metadata.version('xgboost')
        safe_print(f"\nXGBoost 版本: {xgb_version}")
        try:
            import xgboost as xgb
            try:
                params = {'tree_method': 'gpu_hist'}
                xgb.train(params, xgb.DMatrix([[1]], label=[0]), num_boost_round=1)
                safe_print("XGBoost GPU支持: 可用")
            except Exception as e:
                safe_print(f"XGBoost GPU支持: 不可用 ({str(e)})")
        except Exception as e:
            safe_print(f"导入 xgboost 失败: {e}")
    except importlib.metadata.PackageNotFoundError:
        safe_print("\nXGBoost 未安装")
    except Exception as e:
        safe_print(f"检测 XGBoost 时发生异常: {e}")

def check_tensorflow():
    try:
        tf_version = importlib.metadata.version('tensorflow')
        safe_print(f"\nTensorFlow 版本: {tf_version}")
        try:
            gpu_devices = tf.config.list_physical_devices('GPU')
            safe_print(f"GPU 设备: {len(gpu_devices)}")
            if gpu_devices:
                safe_print("GPU 支持: 可用")
                safe_print(f"当前 GPU: {tf.test.gpu_device_name() or 'N/A'}")
                try:
                    with tf.device('/GPU:0'):
                        a = tf.constant([1.0, 2.0, 3.0, 4.0, 5.0, 6.0], shape=[2, 3])
                        b = tf.constant([1.0, 2.0, 3.0, 4.0, 5.0, 6.0], shape=[3, 2])
                        c = tf.matmul(a, b)
                    safe_print("GPU 计算测试: 通过")
                except Exception as e:
                    safe_print(f"GPU 计算测试: 失败 ({str(e)})")
            else:
                safe_print("GPU 支持: 不可用")
                safe_print("可用设备:")
                for device in tf.config.list_physical_devices():
                    safe_print(f"- {device.name} ({device.device_type})")
        except Exception as e:
            safe_print(f"检测 TensorFlow GPU 时发生异常: {e}")
    except importlib.metadata.PackageNotFoundError:
        safe_print("\nTensorFlow 未安装")
    except Exception as e:
        safe_print(f"检测 TensorFlow 时发生异常: {e}")

def main():
    try:
        safe_print("===== 环境信息 =====")
        safe_print(f"Python 版本: {platform.python_version()}")
        safe_print(f"Python 路径: {sys.executable}")
        safe_print(f"工作目录: {os.getcwd()}\n")
        
        safe_print("===== 深度学习框架 =====")
        check_pytorch()
        check_tensorflow()
        
        safe_print("\n===== 梯度提升框架 =====")
        check_lightgbm()
        check_catboost()
        check_xgboost()
    except Exception as e:
        safe_print(f"主流程发生异常: {e}")

if __name__ == "__main__":
    main()